#!/usr/bin/env python3
"""
Apifox MCP服务器
"""

import asyncio
import json
import logging
import os
from typing import Any, Dict, Optional
import requests
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    Tool,
    TextContent,
)
from pydantic import BaseModel, Field

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建MCP服务器实例
server = Server("apifox-mcp-server")

class ApiCallParams(BaseModel):
    """API调用参数模型"""
    project_id: Optional[str] = Field(default=None, description="Apifox项目ID（使用环境变量APIFOX_PROJECT_ID）")
    input: str = Field(description="要导入的OpenAPI JSON字符串数据")
    headers: Optional[Dict[str, str]] = Field(default=None, description="HTTP请求头（自动从环境变量APIFOX_TOKEN添加Authorization）")
    timeout: Optional[int] = Field(default=30, description="请求超时时间（秒）")

@server.list_tools()
async def list_tools() -> list[Tool]:
    """
    列出可用的工具
    """
    return [
        Tool(
            name="import_api",
            description="将OpenAPI、Swagger格式的JSON数据导入到指定的Apifox项目中。",
            inputSchema={
                "type": "object",
                "properties": {
                    "project_id": {
                        "type": "string",
                        "description": "Apifox项目ID，使用环境变量APIFOX_PROJECT_ID"
                    },
                    "input": {
                        "type": "string",
                        "description": "要导入的OpenAPI格式的JSON字符串数据（需要转义）"
                    },
                    "headers": {
                        "type": "object",
                        "description": "HTTP请求头，Authorization会自动从环境变量APIFOX_TOKEN添加",
                        "additionalProperties": {
                            "type": "string"
                        }
                    }
                },
                "required": ["input"]
            }
        )
        ]

@server.call_tool()
async def call_tool(name: str, arguments: dict):
    """
    处理工具调用
    """
    if name != "import_api":
        raise ValueError(f"未知的工具: {name}")
    
    try:
        # 验证和解析参数
        params = ApiCallParams(**arguments)
        
        # 从环境变量获取项目ID
        project_id = params.project_id or os.getenv('APIFOX_PROJECT_ID')
        if not project_id:
            error_msg = "项目ID未提供，请设置APIFOX_PROJECT_ID环境变量"
            logger.error(error_msg)
            raise ValueError(error_msg)
        
        # 验证OpenAPI JSON数据格式
        try:
            openapi_data = json.loads(params.input)
        except json.JSONDecodeError as e:
            error_msg = f"OpenAPI JSON数据格式错误: {str(e)}"
            logger.error(error_msg)
            raise ValueError(error_msg)
        
        # 构建Apifox API URL
        url = f"https://api.apifox.com/v1/projects/{project_id}/import-openapi"
        
        # 准备请求头
        headers = params.headers or {}
        if 'Content-Type' not in headers:
            headers['Content-Type'] = 'application/json'
        
        # 从环境变量获取API Token并添加到请求头
        apifox_token = os.getenv('APIFOX_TOKEN')
        if apifox_token and 'Authorization' not in headers:
            headers['Authorization'] = f'Bearer {apifox_token}'
        
        # 检查是否有认证头
        if 'Authorization' not in headers:
            error_msg = "认证信息未提供，请设置APIFOX_TOKEN环境变量或在headers中提供Authorization"
            logger.error(error_msg)
            raise ValueError(error_msg)
        
        # 自动添加Apifox API版本头
        headers['X-Apifox-Api-Version'] = '2024-03-28'
        
        # 准备请求体
        request_body = {"input": params.input}
        
        logger.info(f"调用Apifox导入API: POST {url}")
        logger.info(f"项目ID: {project_id}")
        
        # 发送HTTP请求
        response = requests.post(
            url=url,
            json=request_body,
            headers=headers,
            timeout=params.timeout
        )
        
        # 构建响应结果
        result = {
            "status_code": response.status_code,
            "reason": response.reason,
            "headers": dict(response.headers),
            "url": response.url,
            "elapsed_time": response.elapsed.total_seconds()
        }
        
        # 尝试解析响应内容
        try:
            if response.headers.get('content-type', '').startswith('application/json'):
                result["response_data"] = response.json()
            else:
                result["response_data"] = response.text
        except Exception as e:
            result["response_data"] = f"无法解析响应内容: {str(e)}"
            result["raw_content"] = response.text[:1000]  # 限制内容长度
        
        # 判断请求是否成功
        success = 200 <= response.status_code < 300
        result["success"] = success
        
        if not success:
            logger.warning(f"API调用失败: HTTP {response.status_code} - {response.reason}")
        else:
            logger.info(f"API调用成功: HTTP {response.status_code}")
        
        if not success:
            raise RuntimeError(f"API调用失败: {json.dumps(result, ensure_ascii=False, indent=2)}")
        
        return [TextContent(
            type="text", 
            text=json.dumps(result, ensure_ascii=False, indent=2)
        )]
        
    except requests.exceptions.Timeout:
        error_msg = f"请求超时（超过{params.timeout}秒）"
        logger.error(error_msg)
        raise RuntimeError(error_msg)
    
    except requests.exceptions.ConnectionError as e:
        error_msg = f"连接错误: {str(e)}"
        logger.error(error_msg)
        raise RuntimeError(error_msg)
    
    except requests.exceptions.RequestException as e:
        error_msg = f"请求异常: {str(e)}"
        logger.error(error_msg)
        raise RuntimeError(error_msg)
    
    except Exception as e:
        error_msg = f"工具执行错误: {str(e)}"
        logger.error(error_msg, exc_info=True)
        raise RuntimeError(error_msg)

async def main():
    """启动MCP服务器"""
    logger.info("启动Apifox MCP服务器...")
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            server.create_initialization_options()
        )

if __name__ == "__main__":
    # 在支持的平台上使用uvloop以提高性能
    try:
        import uvloop
        uvloop.install()
    except ImportError:
        pass
    
    asyncio.run(main()) 